#!/usr/bin/env python3
'''
采用进化策略进行训练网路
这里的算法适合奖励没有负数的情况
'''
import gymnasium as gym
import time
import numpy as np

import torch
import torch.nn as nn

from tensorboardX import SummaryWriter


MAX_BATCH_EPISODES = 100 # 训练的游戏轮数
MAX_BATCH_STEPS = 10000 # 训练的游戏步数
NOISE_STD = 0.01 # 噪声的权重
LEARNING_RATE = 0.001 # 学习率


class Net(nn.Module):
    def __init__(self, obs_size, action_size):
        super(Net, self).__init__()
        # 网络的输出是每个动作的概率
        self.net = nn.Sequential(
            nn.Linear(obs_size, 32),
            nn.ReLU(),
            nn.Linear(32, action_size),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        return self.net(x)


def evaluate(env, net):
    '''
    对加入噪声的网络进行评估

    param env: 游戏环境
    param net: 网络

    return: 总奖励，总步数
    '''
    obs, _ = env.reset()
    reward = 0.0
    steps = 0
    while True:
        obs_v = torch.FloatTensor([obs])
        # 获取预测的动作
        act_prob = net(obs_v)
        # 获取最大概率的动作
        acts = act_prob.max(dim=1)[1]
        # 环境执行动作，获取新环境、奖励、是否结束、其他信息
        obs, r, done, trunc, _ = env.step(acts.data.numpy()[0])
        done = done or trunc
        reward += r
        steps += 1
        if done:
            break
    return reward, steps


def sample_noise(net):
    '''
    这个函数的作用
    进行网络参数的随机采样，并记录其正、负噪声值
    之所以要采样正负噪声，是因为在训练的时候，有可能正向调整参数会使得网络性能提升，也有可能负向调整参数会使得网络性能提升
    所以都要进行采样

    param net: 网络
    return: 正向噪声，逆向噪声
    '''
    pos = []
    neg = []
    for p in net.parameters():
        # 创建一个和网络参数大小一样的噪声
        noise_t = torch.from_numpy(np.random.normal(size=p.data.size()).astype(np.float32))
        # 正向噪声和逆向噪声
        pos.append(noise_t)
        neg.append(-noise_t)
    return pos, neg


def eval_with_noise(env, net, noise):
    '''
    这个函数的作用
    将噪声加入到网络中，然后评估网络的性能，拿到奖励，用于后续训练中判断该噪声
    对网络的性能是否有提升进行更新网络

    param env: 游戏环境
    param net: 网络
    noise: 正向噪声和逆向噪声 噪声的维度和网路每层参数一致

    return 总奖励，总步数
    '''
    # 备份网络参数
    old_params = net.state_dict()
    for p, p_n in zip(net.parameters(), noise):
        # 将网络参数和噪声合并在一起进行更新
        # 将噪声与网络的参数相加
        p.data += NOISE_STD * p_n
    r, s = evaluate(env, net)
    # 还原网络参数
    net.load_state_dict(old_params)
    return r, s


def train_step(net, batch_noise, batch_reward, writer, step_idx):
    '''

    param net: 网络
    param batch_noise: 批量的噪声（里面包含每次生成的正向噪声和逆向噪声，维度和奖励一致）
    param batch_reward: 批量的奖励 （因为分别使用正向噪声和逆向噪声评估网络，所以维度和batch_nosise一致）
    param writer: tensorboard writer记录网络的训练过程
    param step_idx: 当前的训练轮数
    '''

    weighted_noise = None # 每次评估网络时，噪声和奖励的乘积之和 这里使用奖励当作权重
    # 为什么这里更新参数会使得网络的性能提升
    # 这里要结合eval_with_noise方差一起看
    # 在这个方法里面，会将噪声临时加入到网络中，然后测试加入噪声后的网络得到的奖励
    # 如果是正向奖励，那么噪声会使得网络的性能提升，那么在这个方法中会将奖励高的噪声添加进网络，增加这个噪声的权重
    # 由于是CartPole-v0，所以奖励没有负值，所以这里并没有排除负值，将所有的噪音和奖励相乘添加进网络
    # 归一化奖励
    norm_reward = np.array(batch_reward)
    norm_reward -= np.mean(norm_reward) # 减去均值
    s = np.std(norm_reward) # 计算方差
    if abs(s) > 1e-6:
        norm_reward /= s

    for noise, reward in zip(batch_noise, norm_reward):
        if weighted_noise is None:
            # 将噪声和奖励相乘，将矩阵中的噪声和奖励相乘后重新回到矩阵中，所以维度不变
            weighted_noise = [reward * p_n for p_n in noise]
        else:
            for w_n, p_n in zip(weighted_noise, noise):
                # 将噪声和奖励相乘后与矩阵中原始值相加重新放回到权重举证的对应的位置，所以维度不变
                w_n += reward * p_n
    m_updates = []
    # 以网络的参数和权重矩阵进行分组计算更新网路参数
    # 这里更新网络参数是直接将计算的到的噪声奖励权重矩阵直接加到网络参数中
    # 是因为cartpole-v0的奖励没有负值，所以这里直接将所有的噪声和奖励相乘后加入到网络中
    # 如果reward越大，那么这个噪声的权重就越大，其添加到网络中的占比也就越大，也就更改了
    # 网络朝着奖励更大的方向更新
    # 所以这里的方法适合奖励没有负数的游戏
    for p, p_update in zip(net.parameters(), weighted_noise):
        # 将计算的到的噪声奖励权重矩阵每一层的参数，除以权重系数，得到更新的大小
        update = p_update / (len(batch_reward) * NOISE_STD)
        # 与学习率相乘，更新网络参数
        p.data += LEARNING_RATE * update
        # 计算L1范数，l1范数是计算向量中各个元素绝对值之和
        m_updates.append(torch.norm(update))
    # 统计更新参数的绝对值平均值大小
    writer.add_scalar("update_l2", np.mean(m_updates), step_idx)


if __name__ == "__main__":
    writer = SummaryWriter(comment="-cartpole-es")
    # 游戏是简单的平衡杆游戏
    env = gym.make("CartPole-v1")

    net = Net(env.observation_space.shape[0], env.action_space.n)
    print(net)

    step_idx = 0 # 记录训练的次数
    while True:
        t_start = time.time()
        batch_noise = []
        batch_reward = [] # 每次评估获取的奖励
        batch_steps = 0 # 统计评估的总步数
        # 执行一定轮数的游戏
        for _ in range(MAX_BATCH_EPISODES):
            noise, neg_noise = sample_noise(net)
            batch_noise.append(noise)
            batch_noise.append(neg_noise)
            # 将获取的正向噪声加入到网络中，然后执行游戏，获取评估的总奖励和总步数
            reward, steps = eval_with_noise(env, net, noise)
            batch_reward.append(reward)
            batch_steps += steps
            # 将获取的逆向噪声加入到网络中，然后执行游戏，获取评估的总奖励和总步数
            reward, steps = eval_with_noise(env, net, neg_noise)
            # 将获取的奖励加入到奖励列表中
            batch_reward.append(reward)
            batch_steps += steps
            # 判断采集的样本是否达到指定的要求
            if batch_steps > MAX_BATCH_STEPS:
                break

        step_idx += 1
        # 计算平均的奖励
        m_reward = np.mean(batch_reward)
        if m_reward > 199:
            print("Solved in %d steps" % step_idx)
            break

        # 开始训练
        train_step(net, batch_noise, batch_reward, writer, step_idx)
        writer.add_scalar("reward_mean", m_reward, step_idx)
        writer.add_scalar("reward_std", np.std(batch_reward), step_idx)
        writer.add_scalar("reward_max", np.max(batch_reward), step_idx)
        writer.add_scalar("batch_episodes", len(batch_reward), step_idx)
        writer.add_scalar("batch_steps", batch_steps, step_idx)
        speed = batch_steps / (time.time() - t_start)
        writer.add_scalar("speed", speed, step_idx)
        print("%d: reward=%.2f, speed=%.2f f/s" % (step_idx, m_reward, speed))

    pass
